import matplotlib.pyplot as plt
import numpy as np
import shapely
from shapely.affinity import translate
from shapely.ops import split

from target_probability_map import TargetProbabilityMap


def shorten_linestring(line: shapely.LineString, dy: float) -> shapely.LineString:
    if not isinstance(line, shapely.LineString):
        raise ValueError("Input must be a LineString")

    total_length = line.length
    shortened = line.interpolate(dy), line.interpolate(total_length - dy)
    if 2 * dy >= total_length:
        # 缩短长度过大，直接返回空线
        return shapely.LineString()

    # 从dy位置开始，到 length - dy 结束，线的方向不变
    shortened = line.interpolate(dy), line.interpolate(total_length - dy)
    return shapely.LineString([shortened[0], shortened[1]])


def extend_line(line: shapely.LineString, extend_length=1e6):
    coords = list(line.coords)
    x0, y0 = coords[0]
    x1, y1 = coords[-1]

    # 方向向量
    dx = x1 - x0
    dy = y1 - y0
    length = np.hypot(dx, dy)
    dx /= length
    dy /= length

    # 延长起点和终点
    new_start = (x0 - dx * extend_length, y0 - dy * extend_length)
    new_end = (x1 + dx * extend_length, y1 + dy * extend_length)

    return shapely.LineString([new_start, new_end])


def move_control(p1, p2, position, angle, lin_vel, ang_vel, time_now, dt):
    if p2 is None:
        return [0, np.inf]
    else:
        if p2['time'] == time_now:
            return [0, np.inf]
        vel = np.linalg.norm((np.array(p2['pos']) - np.array(position))) / (p2['time'] - time_now)
        direction = np.array(p2['pos']) - np.array(position)
        if np.linalg.norm(direction) < 1e-6:
            return [0, np.inf]
        else:
            ang_vel = (np.arctan2(direction[1], direction[0]) - angle) / dt
            action = [10 * (vel - np.linalg.norm(lin_vel)), ang_vel]
            return action


def path_to_control(task, time_now, position, angle, lin_vel, ang_vel, dt):
    # 判断当前时间位于哪段路径
    p1, p2 = None, None
    for i in range(len(task) - 1):
        if i == 0 and time_now <= task[i]['time']:
            p1 = {
                'pos': position,
                'angle': angle,
                'time': time_now
            }
            p2 = task[i]
            break
        elif i == len(task) - 2 and time_now > task[i + 1]['time']:
            p1 = {
                'pos': position,
                'angle': angle,
                'time': time_now
            }
            p2 = {
                'pos': position,
                'angle': angle,
                'time': time_now
            }
            break
        elif (time_now > task[i]['time']) and (time_now <= task[i + 1]['time']):
            p1 = task[i]
            p2 = task[i + 1]
            break
    return move_control(p1, p2, position, angle, lin_vel, ang_vel, time_now, dt)


def gen_scan_path(pos_stt, area: shapely.Polygon, x_interval, y_margin):
    path = [{
        'pos': np.array(pos_stt),
        'angle': None,
        'speed': None,
        'time': None,
    }]
    # 第一条线取一半间距
    x_min, y_min, x_max, y_max = area.bounds
    scan_line = shapely.LineString(((x_min + x_interval / 2, -1e12), (x_min + x_interval / 2, 1e12)))
    first_inter: shapely.LineString = scan_line.intersection(area)
    # 从下向上扫描
    # 计算区域边缘起点
    start_point = min(first_inter.coords, key=lambda p: p[1])
    # 计算方向
    direction = 1
    path.append({
        'pos': np.array((start_point[0], start_point[1]+y_margin)),
        'angle': None,
        'speed': None,
        'time': None,
    })
    # 向下加一个点
    path.append({
        'pos': np.array((start_point[0], start_point[1]+y_margin - 5)),
        'angle': None,
        'speed': None,
        'time': None,
    })

    # y向运动
    line = shapely.LineString([
        (start_point[0], start_point[1] - 1e12),
        (start_point[0], start_point[1] + 1e12)
    ])
    while True:
        # 计算line和area交点
        inter: shapely.LineString = line.intersection(area)
        if not isinstance(inter, shapely.LineString) or inter.is_empty:
            break
        # 两端缩短 y_margin，并按direction排序，添加进path
        shorted_line = shorten_linestring(inter, y_margin)
        pts = sorted(shorted_line.coords, key=lambda p: direction * p[1])
        for pt in pts:
            path.append({
                'pos': np.array(pt),
                'angle': None,
                'speed': None,
                'time': None,
            })
        # 向右平移,改变方向
        line = translate(line, xoff=x_interval)
        direction = -direction

    return path


def rounded_path(points, radius):
    """
    对一系列点构成的折线进行圆角处理
    :param points: list of {'pos': [float, float], 'angle': float, 'time': float}
    :param radius: 圆角半径
    :return: 新的平滑坐标点列表
    """

    def unit_vector(v):
        if np.linalg.norm(v) < 1e-6:
            return np.zeros_like(v)
        else:
            return v / np.linalg.norm(v)

    def p2p_angle(pt1, pt2):
        return np.arctan2(pt2[1]-pt1[1], pt2[0]-pt1[0])

    new_points = [{
        'pos': np.array(points[0]['pos']),
        'angle': p2p_angle(points[0]['pos'], points[1]['pos']),
        'speed': None,
        'time': None
    }]
    for i in range(1, len(points) - 1):
        p0 = np.array(points[i - 1]['pos'])
        p1 = np.array(points[i]['pos'])
        p2 = np.array(points[i + 1]['pos'])

        v1 = unit_vector(p0 - p1)
        v2 = unit_vector(p2 - p1)

        # 计算夹角夹角的一半余弦定理，确定圆弧起点/终点距折点的距离
        angle = np.arccos(np.clip(np.dot(v1, v2), -1.0, 1.0)) / 2
        if np.tan(angle)>0.25:
            offset = radius / np.tan(angle)

            # 保证 offset 不超过两边距离
            d1 = min(offset, np.linalg.norm(p0 - p1) - 1e-6)
            d2 = min(offset, np.linalg.norm(p2 - p1) - 1e-6)
        else:
            d1 = 0
            d2 = 4*radius

        p1a = p1 + v1 * d1  # 入射点
        p1b = p1 + v2 * d2  # 出射点

        new_points.append({
            'pos': np.array(p1a),
            'angle': p2p_angle(p0, p1),
            'speed': None,
            'time': None
        })
        new_points.append({
            'pos': np.array(p1b),
            'angle': p2p_angle(p1, p2),
            'time': None
        })
    new_points.append({
        'pos': np.array(points[-1]['pos']),
        'angle': p2p_angle(points[-2]['pos'], points[-1]['pos']),
        'speed': None,
        'time': None
    })
    return new_points


def plan_explore_area(boundary, recon_line, agent_pos_list, agent_speed, recon_radius, recon_angle, target_speed, dt, require_time, time_now):
    """
    探索区域数据，返回一个二维numpy数组表示区域。
    1表示可探索区域，0表示不可探索区域。
    """
    # 估计移动到目标区域的时间
    boundary_poly = shapely.Polygon(boundary)
    recon_linestring = shapely.LineString(recon_line)
    recon_boundary_points = recon_linestring.intersection(boundary_poly)
    agent_move_dist = []
    for ap in agent_pos_list:
        agent_to_point_max = np.linalg.norm([np.array(ap) - np.array(ra) for ra in recon_boundary_points.coords], axis=1).max()
        agent_move_dist.append(agent_to_point_max)
    move_time = max(agent_move_dist) / (0.95 * agent_speed)

    # 计算侦察区域
    safer_recon_linestring = translate(recon_linestring, xoff=-max(agent_move_dist) / agent_speed * target_speed)
    safer_recon_linestring = extend_line(safer_recon_linestring, extend_length=1e12)  # 延长线段，避免边界问题
    split_area = split(boundary_poly, safer_recon_linestring)
    recon_poly = max(split_area.geoms, key=lambda p: p.centroid.x)  # 取x最大的区域作为侦察区域
    # 区域范围
    min_x, min_y, max_x, max_y = recon_poly.bounds
    width, height = max_x - min_x, max_y - min_y

    # 区域分割
    # 估计每个侦察区域面积
    subarea_area = (require_time - move_time) * (0.95 * agent_speed) * recon_radius * 0.9
    # 计算区域个数
    recon_num = int(np.ceil(recon_poly.area / subarea_area))
    subareas = []
    # 一个区域
    if recon_num == 1:
        subareas.append(recon_poly)
    # 多个区域
    elif recon_num > 1:
        y_lines = [min_y, max_y]
        y_lines = np.linspace(min_y, max_y, recon_num + 1)
        # 分割
        subareas = []
        remain_area = recon_poly
        for line in y_lines[1:-1]:
            split_line = shapely.LineString(([-1e15, line], [1e15, line]))
            parts = split(remain_area, split_line)
            sorted_parts = sorted(parts.geoms, key=lambda p: p.centroid.y)  # 按照x坐标排序
            subareas.append(sorted_parts[0])
            remain_area = sorted_parts[1]
        subareas.append(remain_area)

    # 计算扫描参数
    # x方向间隔 扫描宽度+目标运动
    sub_height = height / recon_num
    x_interval = (2 * recon_radius * np.sin(recon_angle / 2) - sub_height * 2 / agent_speed * target_speed) * agent_speed / (agent_speed + target_speed) * 0.9
    # y方向边距
    y_margin = recon_radius * np.cos(recon_angle / 2) * 0.95
    # 生成扫描路径
    scan_path = []
    for sa, pos_stt in zip(subareas, agent_pos_list[0:recon_num]):
        scan_path.append(gen_scan_path(pos_stt, sa, x_interval, y_margin))

    # 插入圆角转弯点
    scan_path_angle = []
    for sp in scan_path:
        scan_path_angle.append(rounded_path(sp, 0.2))

    # 计算时间
    # 每个飞机计算每段时间，同一段时间取各飞机最大值
    for i in range(len(scan_path_angle[0])):
        if i == 0:
            for sp in scan_path_angle:
                sp[0]['time'] = time_now
        else:
            # 计算时间间隔
            dt_list = [np.linalg.norm(sp[i - 1]['pos'] - sp[i]['pos']) / (0.95 * agent_speed) for sp in scan_path_angle]

            # 取最大时间间隔
            for sp in scan_path_angle:
                sp[i]['time'] = sp[i - 1]['time'] + max(dt_list)
    # 删除第一个点
    for sp in scan_path_angle:
        sp.pop(0)

    # 返回规定形式
    output = [[
        {'position': (sp['pos'][0], sp['pos'][1], sp['angle']), 'arrival_time': sp['time']}
        for sp in spa
    ] for spa in scan_path_angle]

    return output, subareas


def search_scan(target_prob, recon_radius, recon_angle, position, angle, linear_velocity, angular_velocity, max_spd, sim_interval, world_to_grid):
    # 示例动作，实际应用中应根据扫描结果计算
    h, w = target_prob.shape
    y_indices, x_indices = np.indices((h, w))

    total_weight = target_prob.sum()

    x_centroid = (x_indices * target_prob).sum() / total_weight
    y_centroid = (y_indices * target_prob).sum() / total_weight

    target_pos = np.array((x_centroid / world_to_grid, y_centroid / world_to_grid))

    direction = np.array(target_pos) - np.array(position)
    if np.linalg.norm(direction) < 1e-6:
        return [0, np.inf]
    else:
        ang_vel = (np.arctan2(direction[1], direction[0]) - angle) / sim_interval
        action = [max_spd, ang_vel]

    return action


def search_lock(target_pos, forward_dist, position, angle, linear_velocity, angular_velocity, max_spd, sim_interval):
    # 向目标前方forward_dist点运动。距离目标小于forward_dist时不修改目标点
    pt1 = np.array(position)
    pt2 = np.array(target_pos)

    agent_dir = np.array((np.cos(angle), np.sin(angle)))

    if np.linalg.norm(pt1 - pt2) < 1e-6:
        tar_dir = agent_dir
    else:
        tar_dir = (pt2 - pt1) / np.linalg.norm(pt2 - pt1)

    angle_to_tar = np.arctan2(tar_dir[1], tar_dir[0])

    # 远距离朝向目标
    if abs(angle_to_tar - angle) < np.pi / 2:
        target_pos = pt2 + tar_dir * forward_dist
    else:
        # 远距离背离目标，掉头
        if np.linalg.norm(pt1 - pt2) > forward_dist:
            target_pos = pt2 + tar_dir * forward_dist
        else:
            # 近距离背离目标，继续向前
            target_pos = pt1 - tar_dir * forward_dist

    direction = np.array(target_pos) - np.array(position)
    if np.linalg.norm(direction) < 1e-6:
        return [0, np.inf]
    else:
        ang_vel = (np.arctan2(direction[1], direction[0]) - angle) / sim_interval
        action = [max_spd, ang_vel]

    return action


def test():
    # 区域数据
    x1, x2, y1, y2 = 50000, 200000, 50000 * np.tan(np.pi / 3), 540000
    # 地图区域
    map_poly = [
        (x1, 0), (x1 + x2, 0), (2 * x1 + x2, y1), (2 * x1 + x2, y1 + y2),
        (x1 + x2, 2 * y1 + y2), (x1, 2 * y1 + y2), (0, y1 + y2), (0, y1)
    ]
    # 初始位置
    uav_dy = 10
    agent_pos_list = []
    for i in range(7):
        agent_pos_list.append(
            (0, y1 + y2 / 2 + (i - 2) * uav_dy)
        )
    res, _ = plan_explore_area(
        boundary=[
            (x1, 0), (x1 + x2, 0), (2 * x1 + x2, y1), (2 * x1 + x2, y1 + y2),
            (x1 + x2, 2 * y1 + y2), (x1, 2 * y1 + y2), (0, y1 + y2), (0, y1)
        ],
        recon_line=[[x1+x2, 0], [x1+x2, 2*y1+y2]],
        agent_pos_list=agent_pos_list,
        agent_speed=150,
        recon_radius=60000,
        recon_angle=np.pi/3,
        target_speed=10,
        dt=None,
        require_time=4800,
        time_now=0,
    )
    
    print(res)

    # 绘图
    mat_poly_end = [*map_poly, map_poly[0]]
    plt.plot(
        [p[0] for p in mat_poly_end],
        [p[1] for p in mat_poly_end],
    )
    max_time = 0
    for tt in res:
        pos_x = [p['position'][0] for p in tt]
        pos_y = [p['position'][1] for p in tt]
        angle = np.array([p['position'][2] for p in tt])
        time = [p['arrival_time'] for p in tt]
        max_time = max(max_time, max(time))
        plt.plot(pos_x, pos_y, 'o-')
    print(f"max time: {max_time}")
    plt.axis('equal')
    plt.show()

    return res

def boat_lock(target: TargetProbabilityMap, target_recon_radius, lock_radius, position, max_spd):
    # 目标不明确，在周围等待
    if target.time != target.last_detected:
        tar_ori = target.last_detect_position
        tar_radius = (target.time - target.last_detected) * target.max_speed
        if np.linalg.norm(np.array(tar_ori) - np.array(position)) < 1e-6:
            direction = np.array([-1, 0])
        else:
            direction = (np.array(tar_ori) - np.array(position)) / np.linalg.norm(
                (np.array(tar_ori) - np.array(position)))
        pos_exp = np.array(tar_ori) + direction * (tar_radius + lock_radius)
        
        return pos_exp
    
    # 目标明确，在35km锁定
    else:
        lock_dist = (target_recon_radius + lock_radius)/2
        lock_area_width = target_recon_radius - lock_radius
        lock_area_half = lock_area_width / 2
        if np.linalg.norm(np.array(target.position) - np.array(position)) < 1e-6:
            direction = np.array([-1, 0])
        else:
            direction = (np.array(target.position) - np.array(position)) / np.linalg.norm(
                (np.array(target.position) - np.array(position)))
        speed_exp_rel = direction * max_spd/lock_area_half* (np.linalg.norm(np.array(target.position) - np.array(position))-lock_dist)
        speed_exp = np.array(target.velocity) + speed_exp_rel
        angle_exp = np.arctan2(speed_exp[1], speed_exp[0])
        
        return speed_exp
    
def boat_lock_s1(target: TargetProbabilityMap, target_recon_radius, lock_radius, position, max_spd):
    # USV锁定目标————阶段1
    # 目标不明确，在周围等待
    tar_ori = target.last_detect_position
    tar_radius = (target.time - target.last_detected) * target.max_speed
    if np.linalg.norm(np.array(tar_ori) - np.array(position)) < 1e-6:
        direction = np.array([-1, 0])
    else:
        direction = (np.array(tar_ori) - np.array(position)) / np.linalg.norm(
            (np.array(tar_ori) - np.array(position)))
    pos_exp = np.array(tar_ori) + direction * (tar_radius + lock_radius)
    
    return pos_exp

def boat_lock_s2(target: TargetProbabilityMap, target_recon_radius, lock_radius, position, max_spd):
    # USV锁定目标————阶段2
    # 目标明确，在35km锁定
    lock_dist = (target_recon_radius + lock_radius)/2
    lock_area_width = target_recon_radius - lock_radius
    lock_area_half = lock_area_width / 2
    if np.linalg.norm(np.array(target.position) - np.array(position)) < 1e-6:
        direction = np.array([-1, 0])
    else:
        direction = (np.array(target.position) - np.array(position)) / np.linalg.norm(
            (np.array(target.position) - np.array(position)))
    speed_exp_rel = direction * max_spd/lock_area_half* (np.linalg.norm(np.array(target.position) - np.array(position))-lock_dist)
    speed_exp = np.array(target.velocity) + speed_exp_rel
    angle_exp = np.arctan2(speed_exp[1], speed_exp[0])
    
    return speed_exp

def test_boat_lock():
    pass

if __name__ == '__main__':
    test()
